#!/usr/bin/python
import sys, os, re, time
################################################################################
################################################################################
class Params:
   verbose = True
   def __init__( self, mtxpath ):
      
      assert os.path.isdir( mtxpath )
      self.mtxpath = mtxpath
      
      self.formats = [ 'RMA', 'QTS', 'NZD', 'CSR', 'COO' ]
      self.degrees = [ 2, 4, 8 ]
      
      dagHomeEnvName = 'DAGMTX_HOME_DIR'
      
      if not dagHomeEnvName in os.environ:
         print 'Before using ' + sys.argv[0] + ', you need to set the environment variable "' + dagHomeEnvName + '"'
         print 'to point to the home directory for this project.'
         sys.exit( -1 )
      
      self.dagHomeDir = os.environ[ dagHomeEnvName ]
      self.matlabPath = os.path.join( self.dagHomeDir, 'mat' )
      self.mfileName  = os.path.join( self.matlabPath, 'IdxAndNNZCompactionData.m' )
      self.pysHome = os.path.join( self.dagHomeDir, 'pys' )
      self.buildPy = os.path.join( self.pysHome, 'BuildDAGMtx.py' )
      
      if not os.path.isdir( self.matlabPath ): os.makedirs( self.matlabPath )      
      assert os.path.isdir( self.matlabPath )
   

class DagStats:
   def ParseNZDStats( self, comments ):
      nodeSize = 4 * self.deg
      originalTotalDagBytes = FindValue( comments, value = 'originalTotalDagBytes' )
      pathcompTotalDagBytes = FindValue( comments, value = 'pathcompTotalDagBytes' )
      deltaBytes = originalTotalDagBytes - pathcompTotalDagBytes
      # -- TODO -- remove -- print self.mtxFile
      # -- TODO -- remove -- print 'fmt =', self.fmt
      # -- TODO -- remove -- print 'deg =', self.deg
      # -- TODO -- remove -- print 'originalTotalDagBytes =', originalTotalDagBytes
      # -- TODO -- remove -- print 'pathcompTotalDagBytes =', pathcompTotalDagBytes
      assert deltaBytes >= -nodeSize
      idxsArrBytes = FindValueByTag( comments, tag = 'pattMtx vs. colIdxsAndRowPtrs', value = 'arrBytes' )
      idxsDagBytes = FindValueByTag( comments, tag = 'pattMtx vs. colIdxsAndRowPtrs', value = 'dagBytes' )
      idxsDagBytes -= deltaBytes
      self.idxCompaction = float( idxsArrBytes ) / float( idxsDagBytes )
      self.nnzCompaction = FindValueByTag( comments, tag = 'nnzList vs. nnzList', value = 'compaction' )
      self.pathAdvantage = FindValue( comments, value = 'pathcompAdvantage' )
      self.totCompaction = FindValue( comments, value = 'pathcompTotalCompaction' )
   
   def ParseCSRStats( self, comments ):
      self.nnzCompaction = FindValueByTag( comments, tag = 'nnzList vs. nnzList', value = 'compaction' )
      self.totCompaction = FindValue( comments, value = 'totalCompaction' )
      self.colIdxArrByte = FindValueByTag( comments, tag = 'colIdxs vs. colIdxs', value = 'arrBytes' )
      self.colIdxDagByte = FindValueByTag( comments, tag = 'colIdxs vs. colIdxs', value = 'dagBytes' )
      self.rowPtrArrByte = FindValueByTag( comments, tag = 'rowPtrs vs. rowPtrs', value = 'arrBytes' )
      self.rowPtrDagByte = FindValueByTag( comments, tag = 'rowPtrs vs. rowPtrs', value = 'dagBytes' )
      self.idxArrTotByte = self.colIdxArrByte + self.rowPtrArrByte
      self.idxDagTotByte = self.colIdxDagByte + self.rowPtrDagByte
      self.idxCompaction = float( self.idxArrTotByte ) / float( self.idxDagTotByte )
   
   def ParseCOOStats( self, comments ):
      self.idxCompaction = FindValueByTag( comments, tag = 'pattMtx vs. colIdxsAndRowPtrs', value = 'compaction' )
      self.nnzCompaction = FindValueByTag( comments, tag = 'nnzList vs. nnzList',           value = 'compaction' )
      self.totCompaction = FindValue( comments, value = 'totalCompaction' )
   
   def ReadInStats( self ):
      lines = HeadNLines( self.dagfn, 500 )
      self.valid = True
      self.m   = int( lines[0] )
      self.n   = int( lines[1] )
      self.nnz = int( lines[2] )
      self.numNodes = int( lines[3] )
      self.numBytes = 4 * self.deg * self.numNodes
      self.csrBytes = CSRBytes( self.m, self.n, self.nnz )
      comments = ExtractComments( lines )
      if   self.fmt == 'COO':  self.ParseCOOStats( comments )
      elif self.fmt == 'CSR':  self.ParseCSRStats( comments )
      elif self.fmt == 'NZD':  self.ParseNZDStats( comments )
      else: self.valid = False
      #print self.mtxFile
      #print 'fmt =', self.fmt
      #print 'deg =', self.deg
      #print 'idxCompaction =', self.idxCompaction
      #print 'nnzCompaction =', self.nnzCompaction
      #print 'totCompaction =', self.totCompaction 
      #assert False
   
   def __init__( self, mtxFile, fmt, deg ):
      self.mtxFile = mtxFile
      self.mtxName = MtxFile2Name( mtxFile )
      self.fmt     = fmt
      self.deg     = deg
      self.dagfn   = DagFileName( mtxFile, fmt, deg )
      self.ReadInStats()
   

################################################################################
################################################################################
def AddRedirectionToCmd( cmd, fn ):
  return " ".join([ cmd, "1>", fn, "2>&1" ])

def PrintTimeStamp():
   theTime = time.localtime()
   yr  = '%04d' % theTime.tm_year
   mon = '%02d' % theTime.tm_mon
   day = '%02d' % theTime.tm_mday
   hr  = '%02d' % theTime.tm_hour
   mt  = '%02d' % theTime.tm_min
   sec = '%02d' % theTime.tm_sec
   print '...', '/'.join( [yr, mon, day] ), ':'.join( [hr, mt, sec] )
   sys.stdout.flush()

def SystemWrapperPipelined( cmd, verbose, trial = False ):
   if verbose:
      PrintTimeStamp()
      print '... os.popen(', cmd, ')'
      sys.stdout.flush()
   if not trial:
      FH = os.popen( cmd )
      output = [ line.rstrip() for line in FH.readlines() ]
      FH.close()
      return output
   return []

def SystemWrapper( cmd, verbose, trial = False ):
   dt = 0
   if verbose:
      PrintTimeStamp()
      print '... os.system(', cmd, ')'
      sys.stdout.flush()
   if trial:
      status = 0
   else:
      t0 = time.time()
      status = os.system( cmd )
      t1 = time.time()
      dt = t1 - t0
   assert status == 0
   if verbose:
      print '...', dt, 'seconds elapsed during system call.'
   return dt

def RemoveWrapper( fn, verbose, trial = False ):
	if verbose:
		PrintTimeStamp()
		print '... os.remove(', fn, ')'
		sys.stdout.flush()
	if not trial:
		os.remove( fn )

################################################################################
################################################################################
def WriteLinesToFile( lines, fn ):
   FH = open( fn, 'w' )
   FH.write( '\n'.join( lines ) + '\n' )
   FH.close()

################################################################################
################################################################################
def IntOrFloat( x ):
   try:
      v = int( x )
   except:
      v = float( x )
   return v

def ExtractComments( lines ):
   rawComments = []
   comments = []
   for lineIdx, line in enumerate( lines ):
      if line[0] == '#':
         rawComments.append( line )
      if line == '# begin nodes':
         break
   assert ( 1+lineIdx ) < len( lines )
   for c in rawComments:
      assert c[0] == '#'
      c = c[ 1: ]
      comments.append( c.lstrip().rstrip( 'x' ) )
   return comments

def FindValueByTag( comments, tag = None, value = None ):
   if tag   == None: return None
   if value == None: return None
   tagfound = False
   found = False
   for c in comments:
      if tagfound:
         toks = c.split( ':' )
         if toks[0] == value:
            assert len( toks ) == 2
            found = True
            v = IntOrFloat( toks[1] )
            break
      if c == tag:
         tagfound = True
   assert tagfound
   assert found
   return v

def FindValue( comments, value = None ):
   if value == None: return None
   found = False
   for c in comments:
      toks = c.split( ':' )
      if len( toks ) == 2:
         if toks[0] == value:
            assert not found
            found = True
            v = IntOrFloat( toks[1] )
   try:
      assert found
   except:
      print 'value =', value
      raise
   return v

################################################################################
################################################################################
def HeadNLines( fn, n ):
   cmd = ' '.join([ 'head', '-%d' % n, fn ])
   return SystemWrapperPipelined( cmd, False, False )

def IsRealMtx( mtxFile ):
	line = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
	toks = line.split()
	if len( toks ) != 5:          return False
	if toks[0] != 'MatrixMarket': return False
	if toks[1] != 'matrix':       return False
	if toks[2] != 'coordinate':   return False
	if toks[3] != 'real':         return False
	if toks[4] not in [ 'general', 'symmetric' ]: return False
	return True

def MtxDims( mtxFile ):
   lines = HeadNLines( mtxFile, 200 )
   for line in lines:
      if len( line ):
         if line[0] != '%': break
   toks = line.split()
   assert len( toks ) == 3
   m   = int( toks[0] )
   n   = int( toks[1] )
   nnz = int( toks[2] )
   return ( m, n, nnz )

def CSRBytes( m, n, nnz ):
   nnzValBytes = 8 * nnz
   colIdxBytes = 4 * nnz
   rowPtrBytes = 4 * m + 4
   totalBytes = nnzValBytes + colIdxBytes + rowPtrBytes
   return totalBytes

################################################################################
################################################################################
def MtxFile2Name( mtxFile ):
   ( path, fn ) = os.path.split( mtxFile )
   toks = fn.split( '.' )
   toks.pop()
   rawf = '.'.join( toks )
   return rawf

def DagFileName( mtxFile, fmt, deg ):
   parts = 1
   ( path, fn ) = os.path.split( mtxFile )
   toks = fn.split( '.' )
   toks.pop()
   rawf = '.'.join( toks )
   rawdag = '%s.%s.deg%02d.parts%02d.dag' % ( rawf, fmt.lower(), deg, parts )
   dagfn  = os.path.join( path, rawdag )
   return dagfn

def DagFileExists( mtxFile, fmt, deg ):
   dagfn = DagFileName( mtxFile, fmt, deg )
   if os.path.isfile( dagfn ): return True
   return False

def AllDagsExist( params, mtxFile ):
   for fmt in params.formats:
      for deg in params.degrees:
         if not DagFileExists( mtxFile, fmt, deg ):
            return False
   return True

def SortAndSizeMtxs( mtxFiles ):
   tups = []
   for mtxFile in mtxFiles:
      ( m, n, nnz ) = MtxDims( mtxFile )
      csrBytes = CSRBytes( m, n, nnz )
      tup = ( csrBytes, mtxFile, m, n, nnz )
      tups.append( tup )
   tups.sort()
   return tups

def WalkAndFindMtxs( params ):
   mtxFiles = []
   for ( dirpath, dirs, files ) in os.walk( params.mtxpath ):
      for fn in files:
         toks = fn.split( '.' )
         if toks[ -1 ] == 'mtx':
            mtxFile = os.path.join( dirpath, fn )
            if IsRealMtx( mtxFile ):
               mtxFiles.append( mtxFile )
   return mtxFiles

################################################################################
################################################################################
def GatherStats( params, mtxFileTups ):
   allStats = []
   for ( csrBytes, mtxFile, m, n, nnz ) in mtxFileTups:
      if AllDagsExist( params, mtxFile ):
         bestStats = None
         bestCompaction = 0.0
         for fmt in params.formats:
            for deg in params.degrees:
               stats = DagStats( mtxFile, fmt, deg )
               if stats.valid:
                  if stats.totCompaction > bestCompaction:
                     bestCompaction = stats.totCompaction
                     bestStats = stats
         allStats.append( bestStats )
   return allStats

def WriteMFiles( params, mtxFileTups ):
   statsVec = GatherStats( params, mtxFileTups )
   
   csrBytes   = []
   totCompact = []
   idxCompact = []
   nnzCompact = []
   for stats in statsVec:
      print '%15s, fmt=%3s, deg=%d, tot=%8.2fx, nnz=%8.2fx, idx=%8.2fx' % ( stats.mtxName, stats.fmt, stats.deg, stats.totCompaction, stats.nnzCompaction, stats.idxCompaction )
      csrBytes.append( stats.csrBytes )
      totCompact.append( stats.totCompaction )
      idxCompact.append( stats.idxCompaction )
      nnzCompact.append( stats.nnzCompaction )
   mlines = []
   mlines.append( 'csrBytes   = [' + ','.join([ str(x) for x in csrBytes   ]) + '];' )
   mlines.append( 'totCompact = [' + ','.join([ str(x) for x in totCompact ]) + '];' )
   mlines.append( 'idxCompact = [' + ','.join([ str(x) for x in idxCompact ]) + '];' )
   mlines.append( 'nnzCompact = [' + ','.join([ str(x) for x in nnzCompact ]) + '];' )
   WriteLinesToFile( mlines, params.mfileName )

################################################################################
################################################################################
if len( sys.argv ) < 2:
   ( progPath, progName ) = os.path.split( sys.argv[0] )
   print 'Usage:', progName, '<pathToMtxs>'
   sys.exit( -1 )

params = Params( sys.argv[1] )

mtxFiles    = WalkAndFindMtxs( params )
mtxFileTups = SortAndSizeMtxs( mtxFiles )
WriteMFiles( params, mtxFileTups )
